#!/bin/bash

set -o errexit
set -o xtrace

test_dir=$(realpath $(dirname $0))
. ${test_dir}/../functions

CLUSTER='upgrade'
CLUSTER_SIZE=3
TARGET_API="${API}"
TARGET_OPERATOR_VER=$(echo -n "${API//psmdb.percona.com\/v}" | $sed 's/-/./g')
TARGET_IMAGE="${IMAGE}"
TARGET_IMAGE_MONGOD="${IMAGE_MONGOD}"
TARGET_IMAGE_PMM="${IMAGE_PMM}"
TARGET_IMAGE_BACKUP="${IMAGE_BACKUP}"
if [[ "${TARGET_IMAGE_MONGOD}" == *"percona-server-mongodb-operator"* ]]; then
    MONGO_VER=$(echo -n "${TARGET_IMAGE_MONGOD}" | $sed -r 's/.*([0-9].[0-9])$/\1/')
else
    MONGO_VER=$(echo -n "${TARGET_IMAGE_MONGOD}" | $sed -r 's/.*:([0-9]+\.[0-9]+).*$/\1/')
fi

INIT_OPERATOR_VER=$(curl -s https://check.percona.com/versions/v1/psmdb-operator | jq -r '.versions[].operator' | sort -V | tail -n1)
# if testing on release branch and version service is already updated with new operator images
# use the older version of operator as initial point for test
if [[ "${INIT_OPERATOR_VER}" == "${TARGET_OPERATOR_VER}" ]]; then
    INIT_OPERATOR_VER=$(curl -s https://check.percona.com/versions/v1/psmdb-operator | jq -r '.versions[].operator' | sort -V | tail -n2 | head -n1)
fi
GIT_TAG="v${INIT_OPERATOR_VER}"
INIT_OPERATOR_IMAGES=$(curl -s "https://check.percona.com/versions/v1/psmdb-operator/${INIT_OPERATOR_VER}/latest?databaseVersion=${MONGO_VER}")
OPERATOR_NAME='percona-server-mongodb-operator'
API="psmdb.percona.com/v${INIT_OPERATOR_VER//./-}"
IMAGE=$(echo "${INIT_OPERATOR_IMAGES}" | jq -r '.versions[].matrix.operator[].imagePath')
# we use the starting image from the same repo so we don't need to use initImage option
if [[ "$(echo ${TARGET_IMAGE}|cut -d'/' -f1)" == "perconalab" ]]; then
    IMAGE="${IMAGE/percona\//perconalab\/}"
fi
IMAGE_MONGOD=$(echo "${INIT_OPERATOR_IMAGES}" | jq -r '.versions[].matrix.mongod[].imagePath')
IMAGE_PMM=$(echo "${INIT_OPERATOR_IMAGES}" | jq -r '.versions[].matrix.pmm[].imagePath')
IMAGE_BACKUP=$(echo "${INIT_OPERATOR_IMAGES}" | jq -r '.versions[].matrix.backup[].imagePath')

if [[ "${TARGET_API}" == "${API}" ]]; then
    echo "API and TARGET_API variables are the same: ${API}! Something is wrong!"
    exit 1
fi

function compare_generation() {
    local generation="$1"
    local resource="${2:-statefulset}"
    local name="$3"
    local current_generation

    current_generation=$(kubectl_bin get "${resource}" "${name}" -o jsonpath='{.metadata.generation}')
    if [[ "${generation}" != "${current_generation}" ]]; then
        echo "Generation for resource type ${resource} with name ${name} is: ${current_generation}, but should be: ${generation}!"
        exit 1
    fi
}

function wait_cluster_consistency() {
    until [[ "$(kubectl_bin get psmdb "${CLUSTER}" -o jsonpath='{.status.state}')" == "ready" \
             && "$(kubectl_bin get psmdb "${CLUSTER}" -o jsonpath='{.status.replsets.rs0.ready}')" == "${CLUSTER_SIZE}" ]]; do
        echo 'waiting for cluster readyness'
        sleep 20
    done
}

function check_applied_images() {
    local updated_image="$1"

    case "${updated_image}" in
        "operator")
            if [[ "${TARGET_IMAGE}" == $(kubectl_bin get pod --selector=name="${OPERATOR_NAME}" -o jsonpath='{.items[*].spec.containers[?(@.name == "'"${OPERATOR_NAME}"'")].image}') \
                && "${IMAGE_BACKUP}" == $(kubectl_bin get psmdb "${CLUSTER}" -o jsonpath='{.spec.backup.image}') \
                && "${IMAGE_PMM}" == $(kubectl_bin get psmdb "${CLUSTER}" -o jsonpath='{.spec.pmm.image}') \
                && "${IMAGE_MONGOD}" == $(kubectl_bin get psmdb "${CLUSTER}" -o jsonpath='{.spec.image}') ]]; then
                : Operator image has been updated correctly
            else
                echo 'Operator image has not been updated'
                exit 1
            fi
            ;;
        "all")
            if [[ "${TARGET_IMAGE}" == $(kubectl_bin get pod --selector=name="${OPERATOR_NAME}" -o jsonpath='{.items[*].spec.containers[?(@.name == "'"${OPERATOR_NAME}"'")].image}') \
                && "${TARGET_IMAGE_BACKUP}" == $(kubectl_bin get psmdb "${CLUSTER}" -o jsonpath='{.spec.backup.image}') \
                && "${TARGET_IMAGE_PMM}" == $(kubectl_bin get psmdb "${CLUSTER}" -o jsonpath='{.spec.pmm.image}') \
                && "${TARGET_IMAGE_MONGOD}" == $(kubectl_bin get psmdb "${CLUSTER}" -o jsonpath='{.spec.image}') ]]; then
                : Cluster images have been updated correctly
            else
                echo 'Cluster images have not been updated'
                exit 1
            fi
            ;;
    esac
}

function deploy_operator_gh() {
    desc 'start operator'
    kubectl_bin apply -f "https://raw.githubusercontent.com/percona/percona-server-mongodb-operator/${GIT_TAG}/deploy/crd.yaml"
    kubectl_bin apply -f "https://raw.githubusercontent.com/percona/percona-server-mongodb-operator/${GIT_TAG}/deploy/rbac.yaml"
    curl -s "https://raw.githubusercontent.com/percona/percona-server-mongodb-operator/${GIT_TAG}/deploy/operator.yaml" > "${tmp_dir}/operator_${GIT_TAG}.yaml"
    $sed -i -e "s^image: .*^image: ${IMAGE}^" "${tmp_dir}/operator_${GIT_TAG}.yaml"
    kubectl_bin apply -f "${tmp_dir}/operator_${GIT_TAG}.yaml"

    sleep 2
    wait_pod "$(get_operator_pod)"
}

function prepare_cr_yaml() {
    local cr_yaml="$1"

    # spinup function expects images to have suffix like "-backup"
    # to replace them with images from environment variables
    curl -s "https://raw.githubusercontent.com/percona/percona-server-mongodb-operator/${GIT_TAG}/deploy/cr.yaml" \
        | yq w - "metadata.name" "${CLUSTER}" \
        | yq w - "spec.upgradeOptions.apply" "disabled" \
        | yq w - "spec.replsets[*].size" "${CLUSTER_SIZE}" \
        | yq w - "spec.replsets[*].arbiter.enabled" "false" \
        | yq w - "spec.backup.enabled" "false" \
        | yq w - "spec.image" "" \
        | yq w - "spec.backup.image" -- "-backup" \
        | yq w - "spec.pmm.image" -- "-pmm" > "${cr_yaml}"
}

function main() {
    kubectl_bin delete -f "${src_dir}/deploy/crd.yaml" || true
    kubectl_bin delete -f "${src_dir}/deploy/rbac.yaml" || true
    create_namespace "${namespace}"
    deploy_operator_gh

    curl -s "https://raw.githubusercontent.com/percona/percona-server-mongodb-operator/${GIT_TAG}/deploy/secrets.yaml" > "${tmp_dir}/secrets.yaml"
    kubectl_bin apply -f "${conf_dir}/client.yml" \
                      -f "${tmp_dir}/secrets.yaml"

    local cr_yaml="${tmp_dir}/cr_${GIT_TAG}.yaml"
    prepare_cr_yaml "${cr_yaml}"
    spinup_psmdb "${CLUSTER}-rs0" "${cr_yaml}"

    compare_generation "1" "statefulset" "${CLUSTER}-rs0"
    compare_generation "1" "psmdb" "${CLUSTER}"

    desc 'upgrade operator'
    kubectl_bin apply -f "${src_dir}/deploy/crd.yaml"
    kubectl_bin apply -f "${src_dir}/deploy/rbac.yaml"
    kubectl_bin patch deployment "${OPERATOR_NAME}" \
        -p'{"spec":{"template":{"spec":{"containers":[{"name":"'"${OPERATOR_NAME}"'","image":"'"${TARGET_IMAGE}"'"}]}}}}'
    kubectl_bin rollout status deployment/"${OPERATOR_NAME}"

    desc 'wait for operator upgrade'
    until [[ $(kubectl_bin get pods --selector=name="${OPERATOR_NAME}" \
                    -o custom-columns='NAME:.metadata.name,IMAGE:.spec.containers[0].image' \
                    | grep -vc 'NAME' | awk '{print $1}') -eq 1 ]]; do
        sleep 5
    done
    sleep 10

    desc 'check images and generation after operator upgrade'
    wait_for_running "${CLUSTER}-rs0" "${CLUSTER_SIZE}"
    check_applied_images "operator"
    compare_generation "1" "statefulset" "${CLUSTER}-rs0"
    compare_generation "1" "psmdb" "${CLUSTER}"

    desc 'patch psmdb images and upgrade'
    kubectl_bin patch psmdb "${CLUSTER}" --type=merge --patch '{
        "spec": {
            "crVersion": "'"${TARGET_OPERATOR_VER}"'",
            "image": "'"${TARGET_IMAGE_MONGOD}"'",
            "pmm": { "image": "'"${TARGET_IMAGE_PMM}"'" },
            "backup": { "image": "'"${TARGET_IMAGE_BACKUP}"'" }
        }}'
    sleep 10

    desc 'check cluster after full upgrade'
    wait_for_running "${CLUSTER}-rs0" "${CLUSTER_SIZE}"
    wait_cluster_consistency "${CLUSTER}" "${CLUSTER_SIZE}"
    simple_data_check "${CLUSTER}-rs0" "${CLUSTER_SIZE}"
    check_applied_images "all"
    compare_generation "2" "statefulset" "${CLUSTER}-rs0"
    compare_generation "2" "psmdb" "${CLUSTER}"

    desc 'cleanup'
    kubectl_bin delete -f "${src_dir}/deploy/crd.yaml" || :
    kubectl_bin delete -f "${src_dir}/deploy/rbac.yaml" || :
    kubectl_bin delete pvc --all
    destroy "${namespace}"
}

main
